# -*- coding:UTF-8 -*-
"""
@author: Louis_Yang
@contact: 302810009@qq.com
@time: 2022/3/30 20:41
@file: particle_filter.py
@desc: 实现粒子滤波算法
"""

import numpy as np
import cv2 as cv
import operator

# --------------------------------------#
#       常量定义，类似C中的#define
# --------------------------------------#
PARTICLE_NUMBER = 300  # 定义每个目标产生的粒子数
TRANS_X_STD = 1.05  # 坐标x高斯抽样的方差
TRANS_Y_STD = 0.66  # 坐标y高斯抽样的方差
TRANS_S_STD = 0.01  # 粒子宽高缩放尺寸的高斯抽样标准差
A1 = 2.0  # 二阶动态回归函数的计算因子s
A2 = -1.0  # 二阶动态回归函数的计算因子
B0 = 1.000  # 二阶动态回归函数的计算因子
Eucli_para = 1000  # 粒子与原图不得超过的欧氏距离


# ----------------------------------------------------------------------------------#
#                                 定义“粒子类”
#                 实现类似C中的结构体方法， 每个粒子对象都有多个成员属性,
#                 这里不定义源码中图像宽高, 因为YOLO在detection中输出的是目标的
#                 左上角坐标与右下角坐标，方便直接使用OpenCV画矩形, 这样使用右下角
#                 坐标减去左上角坐标就能计算出粒子的宽和高了，还有就是源码中是用粒子
#                 中心坐标作为二阶动态回归的变量，这里我们用的左上角坐标进行计算。
# ----------------------------------------------------------------------------------#
class Particle(object):  # 不写object也行，object是所有类的父类
    class Struct(object):
        def __init__(self, or_left_corner, or_right_corner,
                     now_left_corner, now_right_corner, now_scale,
                     pre_left_corner, pre_right_corner, pre_scale,
                     hist, dhash, weight):
            self.or_left_corner = or_left_corner  # 粒子"原始"左上角坐标
            self.or_right_corner = or_right_corner  # 粒子"原始"右下角角坐标

            self.now_left_corner = now_left_corner  # 粒子"当前"左上角坐标
            self.now_right_corner = now_right_corner  # 粒子"当前"右下角坐标
            self.now_scale = now_scale  # 当前方框缩放比例

            self.pre_left_corner = pre_left_corner  # 粒子"上一帧"左上角坐标
            self.pre_right_corner = pre_right_corner  # 粒子"上一帧"右下角坐标
            self.pre_scale = pre_scale  # 粒子上一帧"方框缩放"比例

            self.hist = hist  # 粒子"当前"特征直方图
            self.dhash = dhash  # 粒子"当前"差异哈希值
            self.weight = weight  # 粒子"当前"权重

    #  该类方法用于初始化粒子时设置初值
    def set_init_val(self, or_left_corner, or_right_corner,
                     now_left_corner, now_right_corner, now_scale,
                     pre_left_corner, pre_right_corner, pre_scale,
                     hist, dhash, weight):
        return self.Struct(or_left_corner, or_right_corner,
                           now_left_corner, now_right_corner, now_scale,
                           pre_left_corner, pre_right_corner, pre_scale,
                           hist, dhash, weight)


# ------------------------------------------------------------------------#
#                     初始化需要执行搜索任务的粒子个数
#                     通常该函数使用在首次检测到新目标时使用
#                     使每个粒子初始化都与目标参数相同
# -------------------------------------------------------------------------#
def init_particle(particle_num, or_left_corner, or_right_corner,
                  now_left_corner, now_right_corner, now_scale,
                  pre_left_corner, pre_right_corner, pre_scale, hist, dhash, weight):
    """
    :particle_num 粒子数量
    :or_left_corner "原始"粒子的左上角坐标
    :or_right_corner "原始"粒子的右上角坐标
    :now_left_corner "当前"粒子的左上角坐标
    :now_right_corner "当前"粒子的左上角坐标
    :now_scale "当前"粒子的方框缩放因子
    :pre_left_corner "上一帧"粒子的左上角坐标
    :pre_right_corner "上一帧"粒子的左上角坐标
    :pre_scale "上一帧"粒子的方框缩放因子
    :hist "当前"粒子的颜色直方图
    :weight "当前"粒子的权重
    """
    particle = Particle()  # 实例化一个粒子父类
    particle_lis = []  # 用于保存初始化后的粒子列表
    for i in range(particle_num):
        particle_lis.append(particle.set_init_val(or_left_corner, or_right_corner,
                                                  now_left_corner, now_right_corner, now_scale,
                                                  pre_left_corner, pre_right_corner, pre_scale, hist, dhash, weight))

    return particle_lis


# ----------------------------------------------------------#
#                     获取图像的三个通道直方图
#                         以列表形式返回
# ----------------------------------------------------------#
def get_hist(image):
    """
    :image 输入图像，格式为numpy array格式
    :return 返回输入图像中三个通道的颜色直方图
    """
    lis = []
    for i in range(3):  # 三个通道都要计算
        # cv.calcHist(输入图像, 获取通道, 掩膜=None表处理整幅图像, 直方图分成多少份, 直方图中各个像素的值范围)
        hist = cv.calcHist([image], [i], None, [256], [0, 256])  # 获取通道颜色直方图数据
        # normalize()函数将图片的值缩放到0-255之间，方便后续使用compareHist()函数对比
        lis.append(cv.normalize(hist, hist, 0, 1, cv.NORM_MINMAX, -1))  # 归一化数据并放入列表当中

    return lis


# ----------------------------------------------------------#
#                       获取目标区域图像
# ----------------------------------------------------------#
def get_object_area(image, y1, x1, y2, x2):
    """
    :image 输入图像，格式为numpy array格式
    :x1 目标左上角x坐标
    :y1 目标左上角y坐标
    :x2 目标右下角x坐标
    :y2 目标右下角y坐标
    :return 返回目标区域图像
    """
    object_area = image[y1:y2, x1:x2]

    return object_area


# ----------------------------------------------------------------------------#
#                      "对比两个图像的颜色直方图相似程度"
#                  使用颜色直方图作为依据，取每个通道相似度的总和
#                       的平均值作为总相似度，返回相似度
# ----------------------------------------------------------------------------#
def calcu_similar_hist(img1_hist_list, img2_hist_list):
    """
    :img1_hist_list 第一个图像的颜色直方图列表
    :img2_hist_list 第二个图像的颜色直方图列表
    :return 返回最终相似度结果
    """
    similarity = 0  # 保存相似度
    for step, data in enumerate(img1_hist_list):  # 三个通道都要对比
        similarity += cv.compareHist(data, img2_hist_list[step], 0)  # 利用compareHist()进行相似度比较
    similarity /= 3.0  # 取三个通道相似度的平均值
    return similarity


# ----------------------------------------------------------------------------#
#                          "差异哈希算法计算哈希值"
#                  从技术层面来讲，它的速度快是因为不需要计算颜色均值。
#                dHash算法有着aHash算法的速度，和非phash常小的误报率。
# ----------------------------------------------------------------------------#
def get_dhash(img, shape=(10, 10)):
    # 缩放10*11
    img = cv.resize(img, (shape[0] + 1, shape[1]))
    # 转换灰度图
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    hash_str = ''
    # 每行前一个像素大于后一个像素为1，相反为0，生成哈希
    for i in range(shape[0]):
        for j in range(shape[1]):
            if gray[i, j] > gray[i, j + 1]:
                hash_str = hash_str + '1'
            else:
                hash_str = hash_str + '0'
    return hash_str


# -----------------------------------------------------------#
#                   "计算两图像的哈希差异"
# -----------------------------------------------------------#
def calcu_similar_dhash(hash1, hash2, shape=(10, 10)):
    n = 0
    # hash长度不同则返回-1代表传参出错
    if len(hash1) != len(hash2):
        return -1
    # 遍历判断
    for i in range(len(hash1)):
        # 相等则n计数+1，n最终为相似度
        if hash1[i] == hash2[i]:
            n = n + 1
    return n / (shape[0] * shape[1])


# -----------------------------------------------------------------------------------------------------------#
#                                 "改变粒子的位置"
#                    使用二阶动态回归函数来自动更新粒子位置状态和其他参数
#                    以某个坐标点(x, y)中的x为例，作者Rob Hess源码中的思路如下：
#     新的x坐标 = A1*(现在的x坐标 - 原始的x坐标) + A2*(上一帧的x坐标 - 原始的x坐标) + B0*(高斯扰动值) + 原始的x坐标
# -----------------------------------------------------------------------------------------------------------#
def particle_transition(old_par, frame_w, frame_h):
    """
    :frame_w: 原一帧图像的宽
    :frame_h: 原一帧图像的高
    :return 返回变化位置等参数后的粒子
    """
    particle = Particle()  # 实例化一个粒子父类
    new_par = particle.set_init_val([], [], [], [], [], [], [], [], [], [], [])  # 新建一个粒子用来存放transition后的粒子

    # 首先计算出该粒子现在的宽和高，用右下角坐标和左上角坐标即可计算
    weight = old_par.now_right_corner[0] - old_par.now_left_corner[0]  # xmax - xmin得到宽度
    height = old_par.now_right_corner[1] - old_par.now_left_corner[1]  # ymax - ymin得到高度

    # 计算偏移后，新的左上角x坐标
    new_lx = A1 * (old_par.now_left_corner[0] - old_par.or_left_corner[0]) + \
             A2 * (old_par.pre_left_corner[0] - old_par.or_left_corner[0]) + \
             B0 * np.random.normal(loc=0, scale=TRANS_X_STD) + old_par.or_left_corner[0]
    new_lx = max(0.0, min(float(frame_w) - 1.0, new_lx))  # 限制粒子不可超出图像范围
    new_lx = int(np.rint(new_lx))  # 取最接近的整数，因为retangle不允许坐标为浮点数

    # 计算偏移后，新的左上角y坐标
    new_ly = A1 * (old_par.now_left_corner[1] - old_par.or_left_corner[1]) + \
             A2 * (old_par.pre_left_corner[1] - old_par.or_left_corner[1]) + \
             B0 * np.random.normal(loc=0, scale=TRANS_Y_STD) + old_par.or_left_corner[1]
    new_ly = max(0.0, min(float(frame_h) - 1.0, new_ly))  # 限制粒子不可超出图像范围
    new_ly = int(np.rint(new_ly))

    # 计算偏移后，方框尺寸缩放比例
    scale = A1 * (old_par.now_scale - 1.0) + \
            A2 * (old_par.pre_scale - 1.0) + \
            B0 * np.random.normal(loc=0, scale=TRANS_S_STD) + 1.0
    scale = max(0.1, scale)  # 限制缩放比例
    scale = old_par.now_scale
    # 更新粒子参数
    new_par.now_left_corner = [new_lx, new_ly]  # 新的左上角坐标
    new_par.now_right_corner = [int(new_lx + np.rint(weight * scale)),
                                int(new_ly + np.rint(height * scale))]  # 新的右下角坐标
    new_par.now_scale = scale  # 新的缩放尺寸
    new_par.weight = 0  # 更新权重

    new_par.pre_left_corner = old_par.now_left_corner  # 上一帧的左上角坐标为 旧粒子的当前左上角坐标
    new_par.pre_right_corner = old_par.now_right_corner  # 同上
    new_par.pre_scale = old_par.now_scale  # 旧粒子的缩放尺寸

    new_par.or_left_corner = old_par.or_left_corner  # 原始坐标不变化，只是为了计算二阶动态回归时用
    new_par.or_right_corner = old_par.or_right_corner
    return new_par


# -----------------------------------------------------#
#                      "权重归一化"
#                 权重归一化，方便处理和计算
# -----------------------------------------------------#
def normalize_weights(particle_lis):
    """
    :particle_lis 粒子列表
    :return 返回归一化权重后的粒子列表
    """
    sum = 0

    for each_par in particle_lis:
        sum += each_par.weight
    for each_par in particle_lis:
        each_par.weight /= sum

    return particle_lis


# ---------------------------------------------#
#                "冒泡排序" 改进版本
# ---------------------------------------------#
def particle_sort(particle_lis):
    """
    :particle_lis 粒子列表
    :return 返回按照从大到小权重排序好的粒子列表
    """
    num = len(particle_lis)  # 计算总长度
    last_exchange_index = 0  # 记录最后一次交换的位置
    sort_border = num - 1  # 无需数组的边界，每次只需要比较到这个点即可

    for i in range(num):
        flage = True
        for j in range(sort_border):
            if particle_lis[j].weight < particle_lis[j + 1].weight:
                temp = particle_lis[j + 1]
                particle_lis[j + 1] = particle_lis[j]
                particle_lis[j] = temp
                flage = False
                last_exchange_index = j
        sort_border = last_exchange_index
        if flage:
            break
    return particle_lis


def calcu_euclidean_distance(coor1, coor2):
    """
    计算坐标欧式距离，次数省略根号，减少机器运算量
    """
    return np.square(coor1[0] - coor2[0]) + np.square(coor1[1] - coor2[1])


# -----------------------------------------------------------------------------------------------------------#
#                          ⭐⭐⭐根据粒子权重进行"重采样"⭐⭐⭐
#                    使用二阶动态回归函数来自动更新粒子位置状态和其他参数
#                    以某个坐标点(x, y)中的x为例，作者Rob Hess源码中的思路如下：
#     现在的x坐标 = A1*(现在的x坐标 - 原始的x坐标) + A2*(上一帧的x坐标 - 原始的x坐标) + B0*(高斯扰动值) + 原始的x坐标
#     这里比较难理解，请认真阅读下:
#         #======================================================================================#
#         # k为操作数，这一步很重要，目的就是根据粒子权重来决定是否保留，高的保留，低的不保留
#         # 先用权重最大的粒子的权重来乘以粒子成员数量，得到应该保留的粒子数量，然后执行保留操作，
#         # 当粒子操作数连1都不到，说明粒子权重过低，不需要保留，np.rint()函数是返回最接近的整数。
#         #
#         # 举个例子:
#         # 某粒子列表共有粒子30个，如果某粒子k的归一化权重为0.03，则操作数为0.03 * 30 = 0.9 > 0.5，
#         # 因其大于0.5，使用np.rint()后，操作数为1，该粒子可以保留，
#         # 但只保留这一个，如果乘出来大于1了，还是只保留一个，这里是和源码不一样的地方。
#         # 具体可以自己举例子计算。
#         #
#         # 我认为Rob Hess作者的源码中在这个地方的处理是有不足的：
#         # 当计算出来的操作数如果大于1.6(是有这种情况的)时，取得的操作数为2，此时会连续保留两个粒子，
#         # 如果刚好最后一个粒子就是操作数连1都不到的，那就跟着一起保留了，它会跟着下一代继续参与运算，我认为这样是不合理的，
#         # 应当在第一轮不满足操作数大于等于1时就应该被剔除，而不是被保留参与下一代运算。
#         #======================================================================================#
# -----------------------------------------------------------------------------------------------------------#
def resample(particle_lis):
    num = len(particle_lis)
    count = 0  # 计算复制粒子的个数
    # 新建一个新的粒子列表
    new_particle_lis = init_particle(num, None, None, None, None, None, None, None, None, None, None, 0)
    # 首先先对粒子列表按照权重大小进行排序
    particle_lis = particle_sort(particle_lis)

    for step, weight in enumerate(particle_lis):
        k = int(np.rint(particle_lis[step].weight * num))
        if k >= 1:  # 大于1表示该粒子应该保留
            new_particle_lis[count] = particle_lis[step]
            count += 1
        else:  # 进到这里代表从这里开始的后续粒子都不需要继续遍历了，操作数肯定都小于1，因为权重越来越小，
            break

    for step, each_par in enumerate(new_particle_lis):
        if new_particle_lis[step].weight > 0:  # 读取到空的权重说明该粒子还没有被赋值
            # 计算欧式距离，去除扰动过大的粒子
            distance = calcu_euclidean_distance(new_particle_lis[0].now_left_corner,
                                                new_particle_lis[step].now_left_corner)
            if distance > Eucli_para:
                new_particle_lis[step] = new_particle_lis[0]
    if count == num:
        return new_particle_lis
    else:  # 进入到这里代表有权重过小的粒子被抛弃了，那么就不断复制旧的列表中权重最大的粒子填补新的列表，直到补齐原来的成员数量
        while count < num:
            new_particle_lis[count] = particle_lis[0]
            count += 1
        return new_particle_lis


if __name__ == '__main__':

    video_path = './video/test_video5.mp4'
    capture = cv.VideoCapture(video_path)
    frame_num = 0  # 记录输出了多少帧图像
    frame_width = 0  # 记录帧图像的宽
    frame_height = 0  # 记录帧图像的高
    parti_lis = []  # 粒子列表
    object_hist = None  # 记录目标颜色直方图
    object_dhash = None  # 记录目标颜色直方图
    show_each_particle = True
    color = (0, 255, 0)
    # ---------------#
    # 手动定义目标位置
    # ---------------#
    x1 = 840
    y1 = 390
    x2 = 985
    y2 = 640

    while True:
        # 开始逐帧读取视频
        ref, frame = capture.read()
        # 如果读取结果为False ，则代表视频已经读取完毕
        if not ref:
            break

        # 第一帧需要做的事，就是记录目标特征，然后初始化设定数量的粒子
        if frame_num == 0:
            # ---------------#
            #   读取第一帧图像
            # ---------------#
            frame_width = frame.shape[1]
            frame_height = frame.shape[0]
            object_area = frame[y1:y2, x1:x2]  # 裁剪目标区域，裁剪顺序为[y0:y1, x0:x1]
            object_hist = get_hist(object_area)  # 获取目标区域的颜色直方图
            object_dhash = get_dhash(object_area)
            # -----------------------------------------------------------#
            #     初始化粒子列表，并赋值每个粒子与原目标具有相同的信息
            # -----------------------------------------------------------#
            parti_lis = init_particle(PARTICLE_NUMBER, [x1, y1], [x2, y2],
                                      [x1, y1], [x2, y2], 1.0,
                                      [x1, y1], [x2, y2], 1.0,
                                      [object_hist], [object_dhash], [0])
            frame = cv.rectangle(frame, [x1, y1], [x2, y2], color=(255, 0, 0), thickness=2)
            cv.namedWindow('test', 0)
            cv.imshow('test', frame)
            cv.waitKey(0)

        else:
            hist_lis = []
            dhash_lis = []
            similar_lis = []
            similar1_lis = []
            similar2_lis = []
            for step, each_particle in enumerate(parti_lis):
                parti_lis[step] = particle_transition(each_particle, frame_width, frame_height)  # 撒播粒子
                # 取粒子图像区域
                par_area = frame[parti_lis[step].now_left_corner[1]:parti_lis[step].now_right_corner[1],
                           parti_lis[step].now_left_corner[0]:parti_lis[step].now_right_corner[0]]

                par_hist = get_hist(par_area)  # 计算粒子颜色直方图
                parti_lis[step].hist = par_hist  # 存储粒子颜色直方图
                hist_lis.append(par_hist)
                similar1 = calcu_similar_hist(object_hist, par_hist)  # 比较原目标的颜色直方图
                similar1_lis.append(similar1)

                par_dhash = get_dhash(par_area)  # 获取差异哈希值
                parti_lis[step].dhash = par_dhash
                dhash_lis.append(par_dhash)
                similar2 = calcu_similar_dhash(object_dhash, par_dhash)  # 比较原目标的颜色直方图
                similar2_lis.append(similar2)

                similar = (similar1 + similar2) / 2.0
                similar_lis.append(similar)
                parti_lis[step].weight = similar  # 相似度即为权重
                # print(max(similar2_lis))
            max_index, max_similar = max(enumerate(similar_lis), key=operator.itemgetter(1))  # 找到最相似的目标
            if max(similar2_lis) >= 0.70 and max(similar1_lis) >= 0.97:
                print(max(similar2_lis))
                object_hist = hist_lis[max_index]  # 获取目标区域的颜色直方图
                object_dhash = dhash_lis[max_index]
            # 所有粒子权重都要做归一化处理
            parti_lis = normalize_weights(parti_lis)
            # 重要性重采样
            parti_lis = resample(parti_lis)
            parti_lis = resample(parti_lis)

        # 按照权重大小排序粒子列表
        parti_lis = particle_sort(parti_lis)

        if show_each_particle:  # 是否需要在图像上展示所有粒子
            for each_par in parti_lis:
                left_coor = each_par.now_left_corner
                right_coor = each_par.now_right_corner
                frame = cv.rectangle(frame, left_coor, right_coor, color=color, thickness=1)

        # 获取最可能的粒子的坐标
        x1 = parti_lis[0].now_left_corner[0]
        y1 = parti_lis[0].now_left_corner[1]
        x2 = parti_lis[0].now_right_corner[0]
        y2 = parti_lis[0].now_right_corner[1]
        # 更新目标颜色直方图
        # object_hist = get_hi st(frame[y1:y2, x1:x2])

        frame = cv.rectangle(frame, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=3)
        cv.namedWindow('test', 0)
        cv.imshow('test', frame)

        # cv.waitKey(0)
        c = cv.waitKey(1) & 0xff
        if c == 27:
            break

        # 执行代码完毕，帧数+1
        frame_num += 1

"""
改进点1.
    使用欧式距离加强粒子的扰动范围限制

改进点2.
"""
